package com.randi.common.utils;

import com.randi.common.core.domain.dto.TreeNode;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;

/**
 * 菜单工具类
 *
 * @author ziye
 */
public class TreeUtils {

    /**
     * 构建树形结构
     *
     * @param nodeList 节点集合
     * @param <T>      泛型
     * @return 树形结构
     */
    public static <T extends TreeNode> List<T> buildTree(List<T> nodeList) {
        return buildTreeCustom(nodeList, TreeNode::getId, TreeNode::getParentId, (node, parent) -> {
            parent.getChildren().add(node);
        });
    }

    /**
     * 构建树形结构
     *
     * @param nodeList  节点集合
     * @param preAction 节点添加操作
     * @param <T>       泛型
     * @return 树形结构
     */
    public static <T extends TreeNode> List<T> buildTree(List<T> nodeList, BiConsumer<T, T> preAction) {
        return buildTreeCustom(nodeList, TreeNode::getId, TreeNode::getParentId, (node, parent) -> {
            preAction.accept(node, parent);
            parent.getChildren().add(node);
        });
    }

    /**
     * 自定义构建树形结构
     * 自定义指定节点和父级节点对应关系
     * 自定义操作：节点与查找节点匹配成功后，自行添加到子节点，若不匹配成功则添加到根节点
     *
     * @param nodeList     节点列表
     * @param nodeKey      节点 key
     * @param findNodeKey  查找节点 key
     * @param customAction 自定义操作
     * @param <T>          泛型
     * @return 树形结构
     */
    public static <T> List<T> buildTreeCustom(List<T> nodeList, Function<T, Object> nodeKey,
                                              Function<T, Object> findNodeKey, BiConsumer<T, T> customAction) {

        if (nodeList.isEmpty()) {
            return Collections.emptyList();
        }

        Map<Object, T> nodeMap = new HashMap<>(Math.min(nodeList.size(), 16));
        List<T> treeList = new ArrayList<>();

        nodeList.forEach(node -> {
            Object key = nodeKey.apply(node);
            Object findKey = findNodeKey.apply(node);
            nodeMap.put(key, node);

            T findNode = nodeMap.get(findKey);
            if (findNode != null) {
                customAction.accept(node, findNode);
            } else {
                treeList.add(node);
            }
        });

        return treeList;
    }

    /**
     * 自定义构建树形结构
     * 自定义指定节点和父级节点对应关系
     * 自定义节点子节点集合
     *
     * @param nodeList     节点列表
     * @param nodeKey      节点 key
     * @param findNodeKey  匹配节点 key
     * @param nodeChildren 子节点集合
     * @param <T>          泛型
     * @return 树形结构
     */
    public static <T> List<T> buildTreeCustom(List<T> nodeList, Function<T, Object> nodeKey, Function<T, Object> findNodeKey,
                                              Function<T, List<T>> nodeChildren) {

        if (nodeList.isEmpty()) {
            return Collections.emptyList();
        }

        // 保存节点
        Map<Object, T> nodeMap = new HashMap<>(Math.min(nodeList.size(), 16));
        List<T> treeList = new ArrayList<>();

        nodeList.forEach(node -> {
            Object key = nodeKey.apply(node);
            Object findKey = findNodeKey.apply(node);
            nodeMap.put(key, node);

            T findNode = nodeMap.get(findKey);
            if (findNode != null) {
                nodeChildren.apply(findNode).add(node); // 添加子节点
            } else {
                treeList.add(node); // 添加根节点
            }
        });

        return treeList;
    }

}
